PR: Refactoring PyOP3 Codegen by SamSJackson · Pull Request #5392 · firedrakeproject/firedrake

SamSJackson · 2026-08-26T14:47:10Z

Refactoring `pyop3/lower/`

PR presents abstraction in the lower/ design to allow for alternative code generation backends.

Principally, loopy.py has been broken into three separate files:

codegen.py - context orchestrator which returns lowered IR back to translation layer
context.py - generic class and high-level traversal of axis trees to generate pyop3 expressions.
(possible that class and traversal has too many responsibilities).
loopy.py -implements backend-specific lowering from pyop3 to respective target IRs

This PR serves as a stepping stone to the introduction of MLIR as a code generation backend.
A new mlir.py module will subclass context.py that will construct MLIR from PyOP3 expressions, using xDSL.

…ence

- Goal: Create an interface to a code generation context (MLIR or Loopy) - Status: Interface created, battling PETSc bug before cleaning more. - Goal: Integrate MLIR for auto-generation - Status: Was working but transitioning to pyop3->mlir pipeline as opposed to pym->mlir. Refactoring process is ongoing.

requirements tracking pyop3 && typing hints

Merged Connor's update with new setup. Current issue with values loaded into the buffers. Occurs even while using old lower/loopy.py Debugging in process.

PyOP3 updates and path traversals reflected in the new lowering structure

Still need to integrate dtype support for all op3 expressions. Removed abstractmethod while testing.

connorjward

Key point on an initial first pass is I would like to limit the scope of what the context objects know about.

connorjward · 2026-08-26T15:05:08Z

+    else:
+        cs_expr = (insn,)
+
+    if compiler_parameters.codegen == "loopy": 


I think I prefer 'backend' to 'codegen'

That's fine, I will switch. Was only wary that backend is used often but agree that it makes more sense.

Terminology is really confusing around here. We have:

compile

codegen

lower

which are all in some sense interchangeable. I really want to use "compile" as the single term for everything but this conflicts with the other sense of "compiler" (i.e. GCC etc).

I think in my main branch I will do the following renaming:

pyop3/lower to pyop3/compile

pyop3/compile.py to pyop3/cc.py

File naming isn't important for this PR, but I think I will make that change soon and you may hit git conflicts.

connorjward · 2026-08-26T15:17:35Z

+        return ctx 
+
+
+# NOTE: Not a big fan of how compile sits in this file. 


Yeah I would expect to put it in codegen.py. This file can just contain the abstract class.

Moved _compile to codegen.py.

Only point of contention - and why it was not already there - is that parse_loop_properly_this_time uses _compile in the context.py file. As such, _compile now has to be imported. This must be a local import, to avoid a circular import.

It feels like a code smell but it does work.
Possible steps for resolving the potential smell would be:

Move _compile function to another separate file (I think this would bloat the directory)

Rewrite parse_loop_properly_this_time

Else, we can leave it for later, it works fine. May just be a micro-optimisation.

parse_loop_properly_this_time (which I think I would like you to rename) is another function that I don't think should belong to the context and ideally could be rewritten in a backend agnostic fashion.

connorjward · 2026-08-26T15:32:47Z

+
+        return indices
+
+    def compile_standalone_function(


When I designed the 'context' class I imagined that it would be entirely independent from any of the pyop3 'language'. It would only see the buffers. For example here I think we could amend the add_function_call to take in more arguments. The visitor for when we hit standalone function can then be

args = [(a.buffer_view, spec.intent) for a, spec in zip(call.arguments, call.argspec, strict=True)] context.add_function_call(call.function.code, args)

The crucial points are:

The context never sees the StandaloneCalledFunction type

The visitor for StandaloneCalledFunction can remain generic for different backends

I have made some notable changes to follow this idea.

_compile_array_assignment and _compile_loop (formerly _parse_loop...) are now generic in codegen.py

_compile_petsc_mat and _compile_exscan moved back to codegen.py.

Context classes now have abstract methods add_petsc_mat, add_exscan for backend-specific additions.

loopy.py should now be working with buffer_views as opposed to pyop3 terminology, where possible.

Things that I am conscious of:

add_petsc_mat has bloated function definition to avoid passing pyop3 language.

compile_exscan can almost be made agnostic but not until I develop an equivalent to pym.substitute for MLIR gen.

Context classes now have abstract methods add_petsc_mat, add_exscan for backend-specific additions.

I would suggest not bothering with this and just leave the 'generic' routines as they are (i.e. loopy specific). I don't think we care about having MLIR for either of them so just do:

def _compile_petsc_mat(..., ctx): if not isinstance(ctx, LoopyCodegenContext): raise NotImplementedError # loopy-specific code

This makes sense (and my life easier). Pushed these changes

- rename compiler option `codegen` -> `backend` - documenting `backend` compiler option - functional variable renaming

- moving dispatch _compile function from context.py -> codegen.py - local import in _parse_loop to avoid circular import - moving dispatch _collect_temporary_shapes from context.py -> transform.py

- codegen holds more pyop3 traversal. - backend-specific functions introduced where necessary notable flaw is bloated function call for add_petsc_mat.

- not implementing methods for MLIR, loopy-only.

connorjward

I think the structure is very close now. I am not striving for perfection given the rough nature of what you started with, but I think we should make sure that we get a decent chunk of the core design right.

connorjward · 2026-08-27T16:44:12Z

+
+        self.add_assignment(lexpr, rexpr)
+
+    def lower_expr(self, expr, iname_maps, loop_indices, *, intent=READ, paths=None) -> pym.Expression:


I think this is a design issue. We need the codegen context to know how to lower expressions, because it differs between backends, but the logic should really be in another class (LoopyExpressionLowerer(ExpressionLowerer) or LoopyExpressionCompiler(AbstractExpressionCompiler)).

For now the easiest thing to do is just have:

def _lower_expr(...): # a free function ... class LoopyCodegenContext(...): ... def lower_expr(self, ...): return _lower_expr(..., ctx=self)

It does feel bloated in loopy.py. I was conscious of making too many classes.

I think I will revisit making another class once I have a working MLIR lowering implementation.
For now, I have just made this a free function, as suggested.

as title

- moving lower_expr to free function - going to revisit this once I have a working MLIR lowering implementation

SamSJackson added 23 commits June 19, 2026 12:21

sample mlir generation and runnable

efe4ccb

demo to satisfy

39930f6

building plan and stepping through loopy.py to see inputs

94e9c15

designing make_kernel function

285db41

incorporating mlir -> loopy flag and building pymbolic to mlir pipeline

a081fa7

MLIR generating stage - but MLIR typing is incorrect, need type infer…

8b4f6d8

…ence

introducing dtypes for type-inference and testing case

f50a0f7

push before connorjward/pyop3 merge update attempt

b573e66

requirements tracking pyop3 && typing hints

Merge branch 'connorjward/pyop3' into SamSJackson/pyop3-mlir

4c5bac4

Merged Connor's update with new setup. Current issue with values loaded into the buffers. Occurs even while using old lower/loopy.py Debugging in process.

Updating lowering paths for new context class

3358de0

PyOP3 updates and path traversals reflected in the new lowering structure

folder for testing with mlir files

5a1f473

bug fixed, delayed dtype integration

134d0cb

Still need to integrate dtype support for all op3 expressions. Removed abstractmethod while testing.

removing debugging python/text files

8f18475

cleaning up refactor, to add docstrings

c569f30

removing testing files from git

28ea6c9

removing offloading demo files from remote

ea0996f

removing mlir class for merge

950ff8e

adding docstrings

6090c48

removing dtype method for future pr

bf47636

removing debug flags

6a18af4

removing unnecessary requirements-build update

e6b37d5

returning requirements.txt

0fcd377

SamSJackson requested a review from connorjward August 26, 2026 14:47

connorjward added the base:main Run this PR using a main (dev) build label Aug 26, 2026

connorjward requested changes Aug 26, 2026

View reviewed changes

SamSJackson added 4 commits August 27, 2026 10:21

resolving simple PR comments

f82365b

- rename compiler option `codegen` -> `backend` - documenting `backend` compiler option - functional variable renaming

moving _compile to codegen

b45bbf7

- moving dispatch _compile function from context.py -> codegen.py - local import in _parse_loop to avoid circular import - moving dispatch _collect_temporary_shapes from context.py -> transform.py

abstracting backend from pyop3 language

4a2d5ab

- codegen holds more pyop3 traversal. - backend-specific functions introduced where necessary notable flaw is bloated function call for add_petsc_mat.

moving petsc_mat and exscan to codegen

af7d47e

- not implementing methods for MLIR, loopy-only.

connorjward requested changes Aug 27, 2026

View reviewed changes

SamSJackson added 2 commits August 28, 2026 10:31

fix typo and move transform

5dba8f5

as title

restructuring lower_expr

ef3a195

- moving lower_expr to free function - going to revisit this once I have a working MLIR lowering implementation

		return ctx


		# NOTE: Not a big fan of how compile sits in this file.


		self.add_assignment(lexpr, rexpr)

		def lower_expr(self, expr, iname_maps, loop_indices, *, intent=READ, paths=None) -> pym.Expression:

Uh oh!

Conversation

SamSJackson commented Aug 26, 2026

Refactoring pyop3/lower/

Uh oh!

connorjward left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

connorjward left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Refactoring `pyop3/lower/`